import os
import numpy as np
import rasterio
from rasterio.transform import from_origin
"""对降水和气温21年数据平均处理，计算平均值并保存为新的 .tif 文件"""
# 定义文件夹路径
rainfall_path = './mean_rain_tif/'
temperature_path = './mean_tem_tif/'
output_path = './mean_rain_tem_outputs/'  # 输出文件夹路径

# 确保输出文件夹存在
os.makedirs(output_path, exist_ok=True)

# 计算降雨数据的平均值
rainfall_sum = None
rainfall_count = 0

for file in sorted(os.listdir(rainfall_path)):
    if file.endswith('.tif') or file.endswith('.tiff'):
        with rasterio.open(os.path.join(rainfall_path, file)) as src:
            data = src.read(1)  # 读取第一个波段
            nodata = src.nodata
            if nodata is not None:
                data = np.where(data == nodata, np.nan, data)  # 将 NoData 值替换为 NaN
            if rainfall_sum is None:
                rainfall_sum = np.zeros_like(data, dtype=np.float64)
            rainfall_sum += np.nan_to_num(data)  # 累加数据（将 NaN 视为 0）
            print('rainfall:{}'.format(np.nan_to_num(data).max()))
            # rainfall_count += (~np.isnan(data)).astype(int)  # 统计有效像素数量
            rainfall_count +=1

rainfall_mean = rainfall_sum / rainfall_count  # 计算平均值
print('rain_fall_mean.max():{}'.format(rainfall_mean.max()))
# 保存降雨平均值为新的 .tif 文件
with rasterio.open(os.path.join(rainfall_path, sorted(os.listdir(rainfall_path))[0])) as src:
    meta = src.meta
    meta.update(dtype='float32', count=1, nodata=np.nan)
    output_file = os.path.join(output_path, 'mean_rainfall.tif')
    with rasterio.open(output_file, 'w', **meta) as dst:
        dst.write(rainfall_mean.astype('float32'), 1)

print(f"降雨平均值已保存到: {output_file}")

# 计算气温数据的平均值
temperature_sum = None
temperature_count = 0

for file in sorted(os.listdir(temperature_path)):
    if file.endswith('.tif') or file.endswith('.tiff'):
        with rasterio.open(os.path.join(temperature_path, file)) as src:
            data = src.read(1)  # 读取第一个波段
            nodata = src.nodata
            if nodata is not None:
                data = np.where(data == nodata, np.nan, data)  # 将 NoData 值替换为 NaN
            if temperature_sum is None:
                temperature_sum = np.zeros_like(data, dtype=np.float64)
            temperature_sum += np.nan_to_num(data)  # 累加数据（将 NaN 视为 0）
            print('tem:{}'.format(data.max()))
            # temperature_count += (~np.isnan(data)).astype(int)  # 统计有效像素数量
            temperature_count += 1

temperature_mean = temperature_sum / temperature_count  # 计算平均值
print('temperature_mean.max():{}'.format(temperature_mean.max()))
# 保存气温平均值为新的 .tif 文件
with rasterio.open(os.path.join(temperature_path, sorted(os.listdir(temperature_path))[0])) as src:
    meta = src.meta
    meta.update(dtype='float32', count=1, nodata=np.nan)
    output_file = os.path.join(output_path, 'mean_temperature.tif')
    with rasterio.open(output_file, 'w', **meta) as dst:
        dst.write(temperature_mean.astype('float32'), 1)

print(f"气温平均值已保存到: {output_file}")